#include <memory.h>
#include <string>
#include <algorithm>
#include <vector>
#include "ukkonen.h"

// касательно epos, vpos, cpos - мы находимся на ребре, ведущем из вершины vpos с первым символом cpos, на позиции epos - во всех ф-циях

int SuffixTree::ord(char x) const { // получить код символа (от 0..ALPH)
	if (x == termChar) return ALPH - 1;
	if ('Z' >= x && x >= 'A')
		return x - 'A';
	else
		return x - 'a' + 'Z' - 'A' + 1;
}

char SuffixTree::chr(int x) { // символ по коду
	if (x >= ALPH / 2) return 'A' + x;
	else return 'a' + (x - 'Z' + 'A' - 1);
}

void SuffixTree::init(int k) { // инициализация вершины суффиксного дерева
	for (int i = 0; i < ALPH; i++)
		stree[k].e[i].to = -1;
	suff[k] = -1;
	stree[k].height = 0;
	stree[k].parent = 0;
}

int SuffixTree::new_vertex(int p) {
	stree.push_back(vertex());
	suff.push_back(-1);
	init(ssz);
	stree[ssz].parent = p;
	return ssz++;
}

int SuffixTree::split(int v, int c, int l) { // делит ребро суфф. дерева
	if (l == 0) return v;
	if (l == stree[v].e[c].r - stree[v].e[c].l) return stree[v].e[c].to; // ребро из верш. v с первым символом c
	int r = new_vertex(v), tp = ord(t[stree[v].e[c].l + l]); // делим в позиции l
	stree[r].e[tp].to = stree[v].e[c].to;
	stree[r].e[tp].r = stree[v].e[c].r;
	stree[r].e[tp].l = stree[v].e[c].l + l;
	stree[v].e[c].r = stree[v].e[c].l + l;
	stree[v].e[c].to = r;
	stree[stree[r].e[tp].to].parent = r;
	stree[r].height = stree[v].height + l;
	return r;
}

int SuffixTree::add_leaf(int v, int c, int l) { // добавление листа
	int r = new_vertex(v);
	stree[v].e[c].to = r;
	stree[v].e[c].l = l;
	stree[v].e[c].r = t.size();
	stree[r].height = t.size() - l - 1 + stree[v].height;
	return r;
}


bool SuffixTree::try_c(int vpos, int epos, int cpos, int c) { // переход из позиции по символу c
	if (epos == 0)
		return stree[vpos].e[c].to != -1;
	else
		return ord(t[stree[vpos].e[cpos].l + epos]) == c;
}

void SuffixTree::jump(int & vpos, int & epos, int & cpos, int l, int r) { // спуск из позиции на строчку с l по r не включая r
	epos = r - l - stree[vpos].height;
	while (stree[vpos].e[cpos].to != -1 &&
		epos >= stree[vpos].e[cpos].r - stree[vpos].e[cpos].l) {
		int len = stree[vpos].e[cpos].r - stree[vpos].e[cpos].l;
		vpos = stree[vpos].e[cpos].to;
		epos -= len;
		cpos = ord(t[stree[vpos].height + l]);
	}
}

SuffixTree::SuffixTree(std::string a) : t(a + termChar) { 
	length = a.length();
	ssz = 1;
	int l = 0, r = 0;
	suff.push_back(0);
	stree.push_back(vertex()); //создание корня
	init(0);
	suff[0] = 0;
	int vpos = 0;
	int cpos = 0;
	int epos = 0;
	int alone = -1;
	while (l < t.size() && r < t.size()) { // алгоритм Укконена
		int c = ord(t[r]); // символ, который добавляем
		while (l <= r && !try_c(vpos, epos, cpos, c)) { 
			int b = split(vpos, cpos, epos); // делаем split и идём на новую позицию по суфф. ссылке
			if (alone != -1) { suff[alone] = b; alone = -1; }
			if (suff[b] == -1) alone = b;
			add_leaf(b, c, r);
			l++;
			vpos = suff[vpos];
			cpos = ord(t[l + stree[vpos].height]);
			if (l <= r)
				jump(vpos, epos, cpos, l, r);
		}
		if (alone != -1) {
			suff[alone] = vpos;
			alone = -1;
		}
		r++; 
		if (l != r)
			jump(vpos, epos, cpos, l, r);
	}
	init_ssizes(0); 
}

int SuffixTree::init_ssizes(int vert) { // инициализация размеров поддеревьев
	int sret = 0;
	bool leaf = true;
	for (int i = 0; i < ALPH; i++) {
		stree[vert].e[i].ssize = 0;
		if (stree[vert].e[i].to == -1)
			continue;
		leaf = false;
		stree[vert].e[i].ssize = init_ssizes(stree[vert].e[i].to) + stree[vert].e[i].r - stree[vert].e[i].l;
		sret += stree[vert].e[i].ssize;
	}
	if (leaf) return -1;
	return sret;
}

std::string SuffixTree::getKth(int k) { // возврат k-й подстроки
	int vpos = 0, cpos, epos = 0;
	std::string res;
	bool swtch = true; // true - когда стоим в вершине
	while (k > 0) {
		if (!swtch) {
			k--; // пропускаем 1 подстроку, когда поднялись по 1 символу на ребре
			res += t[epos + stree[vpos].e[cpos].l];
			epos++;
			if (epos == stree[vpos].e[cpos].r - stree[vpos].e[cpos].l) { // дошли до конца ребра
				epos = 0;
				vpos = stree[vpos].e[cpos].to;
				swtch = true;
			}
		}
		else {
			for (int i = 0; i < ALPH; i++) { // выбираем, в какое ребро идти дальше
				if (stree[vpos].e[i].to != -1) {
					if (k > stree[vpos].e[i].ssize)
						k -= stree[vpos].e[i].ssize;
					else {
						cpos = i;
						swtch = false;
						break;
					}
				}
			}
		}
	}
	return res;
}

SuffixTree::vertex::vertex() {
	for (int i = 0; i < ALPH; i++)
		e[i].to = -1;
}

int SuffixTree::getLength() const {
	return length;
}